#include "PyPacketFilter.h"

PyPacketFilter::PyPacketFilter()
	: m_pyFilter(NULL)
{
	// Initialize python runtime
	Py_Initialize();
}

PyPacketFilter::~PyPacketFilter()
{
	Py_Finalize();
}

bool PyPacketFilter::SetPySourceFilter(const char* pyCode)
{
	// Compile input python source
	SharedPyObjPtr newPyFilter(Py_CompileString(pyCode, "remote", Py_file_input));

	if(newPyFilter == NULL)
		return false;

	m_pyFilter = newPyFilter;
}

bool PyPacketFilter::operator()(const pcap_pkthdr *h, const u_char *bytes, sockaddr_in& localaddr, sockaddr_in& peeraddr)
{
	// Python local and global variables
	SharedPyObjPtr pyLocals(PyDict_New()), pyGlobals(PyDict_New());

	// Build the python bytearray of the whole packet
	SharedPyObjPtr pyPacket(PyByteArray_FromStringAndSize((const char*)bytes, h->caplen));

	// Register the packet to the python local variables
	PyDict_SetItemString(pyLocals, "packet", pyPacket);

	// Register information of the client TCP connection to the python local variables
	SharedPyObjPtr pySockIP(PyInt_FromLong(ntohl(localaddr.sin_addr.s_addr)));
	PyDict_SetItemString(pyLocals, "localIP", pySockIP);
	SharedPyObjPtr pySockPort(PyInt_FromLong(ntohs(localaddr.sin_port)));
	PyDict_SetItemString(pyLocals, "localPort", pySockPort);
	SharedPyObjPtr pyPeerIP(PyInt_FromLong(ntohl(peeraddr.sin_addr.s_addr)));
	PyDict_SetItemString(pyLocals, "peerIP", pyPeerIP);
	SharedPyObjPtr pyPeerPort(PyInt_FromLong(ntohs(peeraddr.sin_port)));
	PyDict_SetItemString(pyLocals, "peerPort", pyPeerPort);

	// Register the python builtins module
	PyDict_SetItemString(pyGlobals, "__builtins__", PyEval_GetBuiltins());

	// Execute the python filter script
	SharedPyObjPtr pyFilterRet(PyEval_EvalCode((PyCodeObject*)m_pyFilter.get(), pyGlobals, pyLocals));

	// Return the result of the python filter
	if(PyDict_GetItemString(pyLocals, "ret") == Py_True)
		return true;
	return false;
}
